# Testing aggregator.
# 23-10-30

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import joblib
import haiku as hk
import mindspore as ms
import mindspore.nn as ms_nn

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# Test for utils
# setting context & hyper parameters
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

from cybertron.readout.aggregator import EdgeSummation, EdgeMean, NodeSummation, NodeMean, get_aggregator
from ms_cybertron.aggregator.node import TensorMean as ms_NodeMean
from ms_cybertron.aggregator.node import TensorSummation as ms_NodeSummation
from ms_cybertron.aggregator.edge import EdgeTensorMean as ms_EdgeMean
from ms_cybertron.aggregator.edge import EdgeTensorSummation as ms_EdgeSummation

# Test for aggregator
EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

print("=================Test Aggregator====================")
# Test for NodeSummation
test_node_input = np.random.uniform(low=0.1, high=1.1, size=(3, 5, 4)) # (B, A, F)
test_node_mask = np.sum(test_node_input, axis=-1) < 3.0 # (B, A)
test_num_atoms = np.sum(test_node_mask, axis=-1) # (B,)

print("Check input:")
print(test_node_input.shape, test_node_mask.shape, test_num_atoms.shape)
print("Test node input:\n", test_node_input)
print("Test node mask:\n", test_node_mask)
print("Test num atoms:\n", test_num_atoms)

print("------------------Test NodeMean------------------")
def node_fn(x, y, z):
    agg = get_aggregator("node_mean", axis=-2, name="node_mean") # type: ignore
    return agg(x, y, z) # type: ignore

node_fn = jax.vmap(node_fn, in_axes=(0, 0, 0))
ts_node_fn = hk.transform(node_fn, apply_rng=True)
params = ts_node_fn.init(rng, test_node_input, test_node_mask, test_num_atoms)
node_out = jit(ts_node_fn.apply)(params, rng, test_node_input, test_node_mask, test_num_atoms)
print(f"Node out: shape of {node_out.shape}, value of \n {node_out}")

ms_node_fn = ms_NodeMean()
ms_node_input = Tensor(test_node_input, dtype=ms.float32)
ms_node_mask = Tensor(test_node_mask, dtype=ms.bool_)
ms_num_atoms = Tensor(test_num_atoms, dtype=ms.float32).reshape(-1, 1)
ms_node_out = ms_node_fn(ms_node_input, ms_node_mask, ms_num_atoms)

print(f"MS Node out: shape of {ms_node_out.shape}, value of \n {ms_node_out}")
node_same = np.allclose(node_out, ms_node_out.asnumpy(), atol=EPSILON)
print(f"Node aggregator Mean is the same {node_same} in accuracy of {EPSILON}.")
node_diff = node_out - ms_node_out.asnumpy()
node_norm = np.linalg.norm(node_diff)
print(f"norm of diff: {node_norm}.")

# print("------------------Test NodeSummation------------------")
# def node_fn(x, y, z):
#     agg = get_aggregator("node_sum", axis=-2, name="node_sum") # type: ignore
#     return agg(x, y, z) # type: ignore

# node_fn = jax.vmap(node_fn, in_axes=(0, 0, 0))
# ts_node_fn = hk.transform(node_fn, apply_rng=True)
# params = ts_node_fn.init(rng, test_node_input, test_node_mask, test_num_atoms)
# node_out = jit(ts_node_fn.apply)(params, rng, test_node_input, test_node_mask, test_num_atoms)
# print(f"Node out: shape of {node_out.shape}, value of \n {node_out}")

# ms_node_fn = ms_NodeSummation()
# ms_node_input = Tensor(test_node_input, dtype=ms.float32)
# ms_node_mask = Tensor(test_node_mask, dtype=ms.bool_)
# ms_num_atoms = Tensor(test_num_atoms, dtype=ms.float32).reshape(-1, 1)
# ms_node_out = ms_node_fn(ms_node_input, ms_node_mask, ms_num_atoms)

# print(f"MS Node out: shape of {ms_node_out.shape}, value of \n {ms_node_out}")
# node_same = np.allclose(node_out, ms_node_out.asnumpy(), atol=EPSILON)
# print(f"Node aggregator Mean is the same {node_same} in accuracy of {EPSILON}.")
# node_diff = node_out - ms_node_out.asnumpy()
# node_norm = np.linalg.norm(node_diff)
# print(f"norm of diff: {node_norm}.")

# Test for Edge
test_edge_input = np.random.uniform(low=0.1, high=1.1, size=(3, 5, 5, 4)) # (B, A, A, F)
test_edge_mask = np.sum(test_edge_input, axis=-1) < 3.0 # (B, A, A)
test_num_edges = np.sum(test_edge_mask, axis=(-2, -1)) # (B,)

test_edge_cutoff = np.random.uniform(low=0.0, high=1.0, size=(3, 5, 5)) # (B, A, A)

print("Check input:")
print(test_edge_input.shape, test_edge_mask.shape, test_num_edges.shape)
print("Test edge input:\n", test_edge_input)
print("Test edge mask:\n", test_edge_mask)
print("Test num edges:\n", test_num_edges)

print("------------------Test EdgeSummation------------------")
def edge_fn(x, y, z, t):
    agg = get_aggregator("edge_sum", axis=(-3, -2), name="edge_sum") # type: ignore
    return agg(x, y, z, t) # type: ignore

edge_fn = jax.vmap(edge_fn, in_axes=(0, 0, 0, 0))
ts_edge_fn = hk.transform(edge_fn, apply_rng=True)
params = ts_edge_fn.init(rng, test_edge_input, test_edge_mask, test_num_edges, test_edge_cutoff)
edge_out = jit(ts_edge_fn.apply)(params, rng, test_edge_input, test_edge_mask, test_num_edges, test_edge_cutoff)
print(f"Edge out: shape of {edge_out.shape}, value of \n {edge_out}")

ms_edge_fn = ms_EdgeSummation(dim=4, axis=(-3, -2))
ms_edge_input = Tensor(test_edge_input, dtype=ms.float32)
ms_edge_mask = Tensor(test_edge_mask, dtype=ms.bool_)
ms_num_edges = Tensor(test_num_edges, dtype=ms.float32).reshape(-1, 1)
ms_edge_cutoff = Tensor(test_edge_cutoff, dtype=ms.float32)
ms_edge_out = ms_edge_fn(edge_vec=ms_edge_input, 
                         edge_mask=ms_edge_mask, 
                         num_edges=ms_num_edges, 
                         edge_cutoff=ms_edge_cutoff,)

print(f"MS Edge out: shape of {ms_edge_out.shape}, value of \n {ms_edge_out}")
edge_same = np.allclose(edge_out, ms_edge_out.asnumpy(), atol=EPSILON)
print(f"Edge aggregator Sum is the same {edge_same} in accuracy of {EPSILON}.")
edge_diff = edge_out - ms_edge_out.asnumpy()
edge_norm = np.linalg.norm(edge_diff)
print(f"norm of diff: {edge_norm}.")

# # Test for EdgeMean
# print("------------------Test EdgeMean------------------")

# def edge_fn(x, y, z, t):
#     agg = get_aggregator("edge_mean", axis=(-3, -2), name="edge_mean") # type: ignore
#     return agg(x, y, z, t) # type: ignore

# edge_fn = jax.vmap(edge_fn, in_axes=(0, 0, 0, 0))
# ts_edge_fn = hk.transform(edge_fn, apply_rng=True)
# params = ts_edge_fn.init(rng, test_edge_input, test_edge_mask, test_num_edges, test_edge_cutoff)
# edge_out = jit(ts_edge_fn.apply)(params, rng, test_edge_input, test_edge_mask, test_num_edges, test_edge_cutoff)
# print(f"Edge out: shape of {edge_out.shape}, value of \n {edge_out}")

# ms_edge_fn = ms_EdgeMean(dim=4, axis=(-3, -2))
# ms_edge_input = Tensor(test_edge_input, dtype=ms.float32)
# ms_edge_mask = Tensor(test_edge_mask, dtype=ms.bool_)
# ms_num_edges = Tensor(test_num_edges, dtype=ms.float32).reshape(-1, 1)
# ms_edge_cutoff = Tensor(test_edge_cutoff, dtype=ms.float32)
# ms_edge_out = ms_edge_fn(edge_vec=ms_edge_input, 
#                          edge_mask=ms_edge_mask, 
#                          num_edges=ms_num_edges, 
#                          edge_cutoff=ms_edge_cutoff,)

# print(f"MS Edge out: shape of {ms_edge_out.shape}, value of \n {ms_edge_out}")
# edge_same = np.allclose(edge_out, ms_edge_out.asnumpy(), atol=EPSILON)
# print(f"Edge aggregator Sum is the same {edge_same} in accuracy of {EPSILON}.")
# edge_diff = edge_out - ms_edge_out.asnumpy()
# edge_norm = np.linalg.norm(edge_diff)
# print(f"norm of diff: {edge_norm}.")


